"""
数据清洗器核心类
提供各种数据清洗方法
"""

from pandas import DataFrame
import pandas as pd
import numpy as np
import os
from typing import List, Optional
from datetime import datetime

class DataCleaner:
    def __init__(self, file_path: str):
        """
        初始化数据清洗器
        
        Args:
            file_path: 数据文件路径
        """
        self.file_path = file_path
        self.df = self._load_data()
        
    def _load_data(self) -> DataFrame:
        """
        根据文件类型加载数据
        
        Returns:
            DataFrame: 加载的数据框
        """
        if not os.path.exists(self.file_path):
            raise FileNotFoundError(f"文件不存在：{self.file_path}")
            
        if self.file_path.endswith('.csv'):
            return pd.read_csv(self.file_path)
        elif self.file_path.endswith(('.xlsx', '.xls')):
            return pd.read_excel(self.file_path)
        else:
            raise ValueError("不支持的文件类型，仅支持CSV和Excel文件")
    
    def _get_numeric_columns(self, columns: Optional[List[str]] = None) -> List[str]:
        """
        获取数值类型的列
        
        Args:
            columns: 指定的列名列表，如果为None则检查所有列
            
        Returns:
            List[str]: 数值类型的列名列表
        """
        if columns is None:
            return list(self.df.select_dtypes(include=[np.number]).columns)
        return [col for col in columns if col in self.df.select_dtypes(include=[np.number]).columns]
    
    def _get_text_columns(self, columns: Optional[List[str]] = None) -> List[str]:
        """
        获取文本类型的列
        
        Args:
            columns: 指定的列名列表，如果为None则检查所有列
            
        Returns:
            List[str]: 文本类型的列名列表
        """
        if columns is None:
            return list(self.df.select_dtypes(include=['object']).columns)
        return [col for col in columns if col in self.df.select_dtypes(include=['object']).columns]
            
    def fill_missing(self, method: str, columns: Optional[List[str]] = None) -> 'DataCleaner':
        """
        填充缺失值
        
        Args:
            method: 填充方法 ('mean', 'median', 'mode')
            columns: 要处理的列名列表，如果为None则处理所有列
            
        Returns:
            self: 支持方法链式调用
        """
        if method == 'mean':
            return self.fill_missing_mean(columns)
        elif method == 'median':
            return self.fill_missing_median(columns)
        elif method == 'mode':
            return self.fill_missing_mode(columns)
        else:
            raise ValueError(f"不支持的填充方法: {method}")
            
    def fill_missing_mean(self, columns: Optional[List[str]] = None) -> 'DataCleaner':
        """
        使用均值填充数值列的缺失值，文本列使用'NOT-FOUND'填充
        
        Args:
            columns: 要处理的列名列表，如果为None则处理所有列
            
        Returns:
            self: 支持方法链式调用
        """
        # 处理数值列
        numeric_columns = self._get_numeric_columns(columns)
        if numeric_columns:
            self.df[numeric_columns] = self.df[numeric_columns].fillna(self.df[numeric_columns].mean())
        
        # 处理文本列
        text_columns = self._get_text_columns(columns)
        if text_columns:
            self.df[text_columns] = self.df[text_columns].fillna('NOT-FOUND')
            
        return self
        
    def fill_missing_median(self, columns: Optional[List[str]] = None) -> 'DataCleaner':
        """
        使用中位数填充数值列的缺失值，文本列使用'NOT-FOUND'填充
        
        Args:
            columns: 要处理的列名列表，如果为None则处理所有列
            
        Returns:
            self: 支持方法链式调用
        """
        # 处理数值列
        numeric_columns = self._get_numeric_columns(columns)
        if numeric_columns:
            self.df[numeric_columns] = self.df[numeric_columns].fillna(self.df[numeric_columns].median())
        
        # 处理文本列
        text_columns = self._get_text_columns(columns)
        if text_columns:
            self.df[text_columns] = self.df[text_columns].fillna('NOT-FOUND')
            
        return self
        
    def fill_missing_mode(self, columns: Optional[List[str]] = None) -> 'DataCleaner':
        """
        使用众数填充数值列的缺失值，文本列使用'NOT-FOUND'填充
        
        Args:
            columns: 要处理的列名列表，如果为None则处理所有列
            
        Returns:
            self: 支持方法链式调用
        """
        if columns is None:
            columns = self.df.columns
            
        for col in columns:
            if col in self.df.columns:
                if col in self._get_numeric_columns([col]):
                    # 处理数值列
                    mode_value = self.df[col].mode()
                    if not mode_value.empty:
                        self.df[col] = self.df[col].fillna(mode_value[0])
                else:
                    # 处理文本列
                    self.df[col] = self.df[col].fillna('NOT-FOUND')
                    
        return self
        
    def save(self, output_path: str) -> str:
        """
        保存清洗后的数据
        
        Args:
            output_path: 输出文件路径
            
        Returns:
            str: 保存的文件路径
        """
        # 确保输出目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 根据文件扩展名保存
        if output_path.endswith('.csv'):
            self.df.to_csv(output_path, index=False)
        elif output_path.endswith(('.xlsx', '.xls')):
            self.df.to_excel(output_path, index=False)
        else:
            raise ValueError("不支持的输出文件类型，仅支持CSV和Excel文件")
            
        return output_path
        
    def get_preview(self, rows: int = 5) -> dict:
        """
        获取数据预览
        
        Args:
            rows: 预览的行数
            
        Returns:
            dict: 包含数据预览信息的字典
        """
        return {
            'head': self.df.head(rows).to_dict('records'),
            'info': {
                'total_rows': len(self.df),
                'total_columns': len(self.df.columns),
                'missing_values': self.df.isnull().sum().to_dict()
            }
        } 